from typing import List, Dict, Any, TypedDict, Annotated

from langgraph.constants import START
from typing_extensions import TypedDict
from langgraph.graph import StateGraph, END
from langgraph.types import Send
import operator

class OverallState(TypedDict):
    subjects: list[str]
    jokes: Annotated[list[str], operator.add]


# 解析主题，返回多个 Send
def parse_subjects(state: OverallState) -> List[Send]:
    return [Send("generate_joke", {"subject": s, "jokes": []}) for s in state["subjects"]]


# 为每个主题生成一个笑话
def generate_joke(state: Dict[str, Any]) -> Dict[str, Any]:
    subject = state["subject"]
    joke = f"Why did the {subject} cross the road? To get to the other side!"
    return {"jokes": state["jokes"] + [joke]}


# 汇总所有笑话
def combine_jokes(state_list: List[Dict[str, Any]]) -> Dict[str, Any]:
    combined_jokes = []
    for state in state_list:
        combined_jokes.extend(state.get("jokes", []))
    return {"jokes": combined_jokes}


if __name__ == '__main__':
    graph = StateGraph(OverallState)

    # ✅ 注册生成笑话和合并笑话节点
    graph.add_node("generate_joke", generate_joke)
    graph.add_node("combine_jokes", combine_jokes)
    # ✅ 添加条件边，不需要 add_node("parse_subjects", ...)
    graph.add_conditional_edges(START, parse_subjects)
    # graph.add_node("parse_subjects", parse_subjects, is_sender=True)  # 👈 正确注册为 Send 类型节点
    graph.add_edge("generate_joke", "combine_jokes")
    graph.add_edge("combine_jokes", END)

    graph.set_entry_point("parse_subjects")

    app = graph.compile()
    output = app.invoke({"subjects": ["cat", "dog", "programmer"]})
    print(output["jokes"])
